from __future__ import division
import cv2
import numpy as np
from numpy import random
import math
from sklearn.utils import shuffle
from copy import deepcopy
from LMdataset import kpPairs

__all__ = ['Compose','ResizeImg', 'RandomHflip', 'Normalize', 'GenGauMask', 'RandomCrop',
           'RandomRotate']

oriIdx =  [0,1, 2, 3,4, 5,6, 7,8, 9,10,11,12, 13,14, 15,16, 17,18, 19, 20,21,22,23]
flipIdx = [1,0, 2, 4,3, 6,5, 8,7, 11,12,9,10, 14,13, 16,15, 18,17, 19, 22,23,20,21]


def scale_down(src_size, size):
    w, h = size
    sw, sh = src_size
    if sh < h:
        w, h = float(w * sh) / h, sh
    if sw < w:
        w, h = sw, float(h * sw) / w
    return int(w), int(h)


def fixed_crop(src, x0, y0, w, h, size=None):
    out = src[y0:y0 + h, x0:x0 + w]
    if size is not None and (w, h) != size:
        out = cv2.resize(out, (size[0], size[1]), interpolation=cv2.INTER_CUBIC)
    return out


def center_crop(src, size):
    h, w = src.shape[0:2]
    new_w, new_h = scale_down((w, h), size)

    x0 = int((w - new_w) / 2)
    y0 = int((h - new_h) / 2)

    out = fixed_crop(src, x0, y0, new_w, new_h, size)
    return out


def bottom_crop(src, size):
    h, w = src.shape[0:2]
    new_w, new_h = scale_down((w, h), size)

    x0 = int((w - new_w) / 2)
    y0 = int((h - new_h) * 0.75)

    out = fixed_crop(src, x0, y0, new_w, new_h, size)
    return out

def CenterLabelHeatMap(img_width, img_height, c_x, c_y,sigma):
    Y1 = np.linspace(0, img_height-1, img_height)
    X1 = np.linspace(0, img_width-1, img_width)
    [X, Y] = np.meshgrid(X1, Y1)
    X = X - c_x
    Y = Y - c_y
    D2 = X * X + Y * Y
    E2 = 2. * sigma * sigma
    Exponet = D2 / E2
    heatmap = np.exp(-Exponet)
    return heatmap

def rotate_nobound(image,kp, angle, center=None, scale=1.):
    (h, w) = image.shape[:2]
    kp_new = kp.copy()
    kp_new[:, 0] = kp_new[:, 0] * w
    kp_new[:, 1] = kp_new[:, 1] * h
    kp_new[:, 2] = 1

    # if the center is None, initialize it as the center of
    # the image
    if center is None:
        center = (w // 2, h // 2)

    # perform the rotation
    M = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, M, (w, h))

    kp_new[:, 0:2] = np.dot(kp_new, M.T)

    kp_new[:, 0] = kp_new[:, 0] / w
    kp_new[:, 1] = kp_new[:, 1] / h
    kp_new[:, 2] = kp[:,2]

    kp_new[(kp_new[:, 0]<0) | (kp_new[:, 1]<0), 2] = -1


    return rotated, kp_new


def rotate_bound(image, kp, angle):
    # grab the dimensions of the image and then determine the
    # center
    h, w = image.shape[:2]

    kp_new = kp.copy()
    kp_new[:, 0] = kp_new[:, 0] * w
    kp_new[:, 1] = kp_new[:, 1] * h
    kp_new[:, 2] = 1

    (cX, cY) = (w // 2, h // 2)

    M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY

    kp_new[:, 0:2] = np.dot(kp_new, M.T)

    kp_new[:, 0] = kp_new[:, 0] / nW
    kp_new[:, 1] = kp_new[:, 1] / nH
    kp_new[:, 2] = kp[:,2]

    kp_new[(kp_new[:, 0]<0) | (kp_new[:, 1]<0), 2] = -1

    mean_value = image.mean(0).mean(0)
    rotated = cv2.warpAffine(image, M, (nW, nH))

    return rotated, kp_new


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, *args):
        for t in self.transforms:
            args = t(*args)
        return args


class ResizeImg(object):
    def __init__(self, size,inter=cv2.INTER_LINEAR):
        self.size = size
        self.inter = inter

    def __call__(self, img, lm_mask, flm, ori_size):
        return cv2.resize(img, (self.size[1], self.size[0]), interpolation=self.inter), \
                lm_mask, \
                flm, \
                ori_size


class RandomHflip(object):
    def __call__(self, img,lm_mask, flm, ori_size):
        if random.randint(2):
            flm_ = deepcopy(flm)
            flm_[flm_[:, 2] > -1, 0] = 1. - flm_[flm_[:, 2] > -1, 0]
            flm = flm_[flipIdx, :]

            return cv2.flip(img, 1),lm_mask, flm, ori_size
        else:
            return img, lm_mask, flm, ori_size


class RandomCrop(object):
    def __init__(self, margin=0.1):
        self.margin = margin

    def __call__(self, img, lm_mask, flm, ori_size):
        h, w = img.shape[0:2]

        # find bound point
        ymin, ymax = flm[flm[:,2]>-1, 1].min(), flm[flm[:,2]>-1, 1].max()
        xmin, xmax = flm[flm[:,2]>-1, 0].min(), flm[flm[:,2]>-1, 0].max()

        ymin, ymax = int(ymin * h), int(ymax * h)
        xmin, xmax = int(xmin * w), int(xmax * w)

        dxmin = np.random.randint(0, max(int(xmin * (1-self.margin)), 1))
        dymin = np.random.randint(0, max(int(ymin * (1-self.margin)), 1))
        dxmax = np.random.randint(0, max(int((w-xmax) * (1-self.margin)), 1))
        dymax = np.random.randint(0, max(int((h-ymax) * (1-self.margin)), 1))

        new_xmin = 0 + dxmin
        new_ymin = 0 + dymin
        new_xmax = w - dxmax - 1
        new_ymax = h - dymax - 1
        new_w = new_xmax - new_xmin
        new_h = new_ymax - new_ymin

        # kp coord changed
        coord_x = (flm[:, 0] * w).astype(int)
        coord_y = (flm[:, 1] * h).astype(int)
        coord_x  = coord_x - new_xmin
        coord_y = coord_y - new_ymin

        flm[:, 0] = coord_x / new_w
        flm[:, 1] = coord_y / new_h

        img = fixed_crop(img, new_xmin, new_ymin, new_w, new_h, size=None)

        return img, lm_mask, flm, ori_size


class RandomRotate(object):
    def __init__(self, angles, bound):
        self.angles = angles
        self.bound = bound

    def __call__(self, img, lm_mask, flm, ori_size):
        angle = np.random.uniform(self.angles[0], self.angles[1])
        if self.bound:
            img, flm = rotate_bound(img, flm, angle)
        else:
            img, flm = rotate_nobound(img, flm, angle)
        return img, lm_mask, flm, ori_size


class GenGauMask(object):
    def __init__(self,r):
        self.r = r

    def __call__(self, img, lm_mask, flm, ori_size):
        h, w = img.shape[0:2]
        mask_h = int(h / self.r)
        mask_w = int(w / self.r)
        lm_mask = np.zeros((flm.shape[0], mask_h, mask_w), dtype=float)
        for i in range(flm.shape[0]):
            x = min(int(round(w * flm[i,0])), w-1)
            y = min(int(round(h * flm[i,1])), h-1)
            hm = CenterLabelHeatMap(w, h, x, y, sigma=7)
            lm_mask[i] = cv2.resize(hm, dsize=(mask_w, mask_h),interpolation=cv2.INTER_LINEAR)

        return img, lm_mask, flm, ori_size



class Normalize(object):
    def __init__(self,mean, std):
        '''
        :param mean: RGB order
        :param std:  RGB order
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        '''
        self.mean = np.array(mean).reshape(3,1,1)
        self.std = np.array(std).reshape(3,1,1)

    def __call__(self, img,lm_mask, flm, ori_size):
        '''
        :param image:  (H,W,3)  RGB
        :return:
        '''
        return (img.transpose((2, 0, 1)) / 255. - self.mean) / self.std, \
               lm_mask, \
               flm, ori_size




if __name__=='__main__':

    from LMdata.LMdataset import LMdata
    import pandas as pd
    import os
    from utils.plotting import DrawKeyPoints
    from matplotlib import pyplot as plt

    from imutils import rotate_bound, rotate

    class trainAug(object):
        def __init__(self):
            self.augment = Compose([
                # RandomCrop(),
                ResizeImg(size=(336,336)),
                RandomHflip(),
                GenGauMask(r=1),
                # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])

        def __call__(self, *args):
            return self.augment(*args)




    img_root = '/media/gserver/data/landmark/rawdata/train'
    annotation = pd.read_csv(os.path.join(img_root, "Annotations/train.csv"))
    annotation['image_id'] = annotation['image_id'].apply(lambda x: os.path.join(img_root, x))



    data_set = {}
    data_set['train'] = LMdata(annotation, trainAug())

    img, lm_mask, vis_mask, ori_size, flm, cate_idx = data_set['train'][175]

    img = img.numpy().astype(np.uint8)
    # lm_mask = lm_mask.numpy()
    #
    # img_kp = DrawKeyPoints(img.copy(), flm)
    #
    # plt.subplot(121)
    # plt.imshow(img)
    # plt.subplot(122)
    # plt.imshow(img_kp)
    # plt.show()

